from app import db
from datetime import datetime

class ImageTask(db.Model):
    """图片生成任务表"""
    __tablename__ = 'image_tasks'
    
    id = db.Column(db.String(36), primary_key=True, comment='任务UUID')
    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False, comment='用户ID')
    style_template_id = db.Column(db.Integer, db.ForeignKey('style_templates.id'), nullable=False, comment='风格模板ID')
    ai_config_id = db.Column(db.Integer, db.ForeignKey('ai_configs.id'), nullable=False, comment='AI配置ID')
    
    # 图片信息
    original_image_url = db.Column(db.String(500), nullable=False, comment='原始图片URL')
    generated_image_url = db.Column(db.String(500), nullable=True, comment='生成图片URL')
    
    # 任务状态
    status = db.Column(db.String(20), default='pending', comment='任务状态: pending, processing, completed, failed')
    progress = db.Column(db.Integer, default=0, comment='处理进度 0-100')
    
    # 提示词信息
    positive_prompt = db.Column(db.Text, nullable=True, comment='正向提示词')
    negative_prompt = db.Column(db.Text, nullable=True, comment='负向提示词')
    custom_prompt = db.Column(db.Text, nullable=True, comment='用户自定义提示词')
    
    # AI参数
    ai_model = db.Column(db.String(100), nullable=True, comment='使用的AI模型')
    width = db.Column(db.Integer, default=512, comment='图片宽度')
    height = db.Column(db.Integer, default=512, comment='图片高度')
    steps = db.Column(db.Integer, default=20, comment='生成步数')
    cfg_scale = db.Column(db.Float, default=7.0, comment='CFG Scale')
    
    # 错误信息
    error_message = db.Column(db.Text, nullable=True, comment='错误信息')
    
    # 时间戳
    created_at = db.Column(db.DateTime, default=datetime.utcnow, comment='创建时间')
    started_at = db.Column(db.DateTime, nullable=True, comment='开始处理时间')
    completed_at = db.Column(db.DateTime, nullable=True, comment='完成时间')
    
    # 关联关系
    style_template = db.relationship('StyleTemplate', backref='tasks')
    ai_config = db.relationship('AIConfig', backref='tasks')
    
    def __repr__(self):
        return f'<ImageTask {self.id}>'
    
    def to_dict(self):
        return {
            'id': self.id,
            'user_id': self.user_id,
            'status': self.status,
            'progress': self.progress,
            'original_image_url': self.original_image_url,
            'generated_image_url': self.generated_image_url,
            'positive_prompt': self.positive_prompt,
            'negative_prompt': self.negative_prompt,
            'custom_prompt': self.custom_prompt,
            'ai_model': self.ai_model,
            'width': self.width,
            'height': self.height,
            'steps': self.steps,
            'cfg_scale': self.cfg_scale,
            'error_message': self.error_message,
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'started_at': self.started_at.isoformat() if self.started_at else None,
            'completed_at': self.completed_at.isoformat() if self.completed_at else None,
            'style_template': self.style_template.to_dict() if self.style_template else None
        }
    
    def update_status(self, status, progress=None, error_message=None):
        """更新任务状态"""
        self.status = status
        if progress is not None:
            self.progress = progress
        if error_message is not None:
            self.error_message = error_message
            
        if status == 'processing' and not self.started_at:
            self.started_at = datetime.utcnow()
        elif status in ['completed', 'failed']:
            self.completed_at = datetime.utcnow()
            
        db.session.commit()